// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/expression.h>

#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <nebula/compute/api_scalar.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/cast_internal.h>
#include <nebula/compute/registry.h>
#include <nebula/core/record_batch.h>
#include <nebula/core/table.h>
#include <turbo/log/logging.h>

namespace nebula::compute {

    using internal::GetCastFunction;

    struct KnownFieldValues {
        std::unordered_map<FieldRef, Datum, FieldRef::Hash> map;
    };

    inline const Expression::Call *CallNotNull(const Expression &expr) {
        auto call = expr.call();
                DKCHECK_NE(call, nullptr);
        return call;
    }

    inline std::vector<TypeHolder> get_types(const std::vector<Expression> &exprs) {
        std::vector<TypeHolder> types(exprs.size());
        for (size_t i = 0; i < exprs.size(); ++i) {
                    DKCHECK(exprs[i].is_bound());
            types[i] = exprs[i].type();
        }
        return types;
    }

    inline std::vector<TypeHolder> get_types(const std::vector<Datum> &values) {
        std::vector<TypeHolder> types(values.size());
        for (size_t i = 0; i < values.size(); ++i) {
            types[i] = values[i].type();
        }
        return types;
    }

    struct Comparison {
        enum type {
            NA = 0,
            EQUAL = 1,
            LESS = 2,
            GREATER = 4,
            NOT_EQUAL = LESS | GREATER,
            LESS_EQUAL = LESS | EQUAL,
            GREATER_EQUAL = GREATER | EQUAL,
        };

        static const type *Get(const std::string &function) {
            static std::unordered_map<std::string, type> map{
                    {"equal",         EQUAL},
                    {"not_equal",     NOT_EQUAL},
                    {"less",          LESS},
                    {"less_equal",    LESS_EQUAL},
                    {"greater",       GREATER},
                    {"greater_equal", GREATER_EQUAL},
            };

            auto it = map.find(function);
            return it != map.end() ? &it->second : nullptr;
        }

        static const type *Get(const Expression &expr) {
            if (auto call = expr.call()) {
                return Comparison::Get(call->function_name);
            }
            return nullptr;
        }

        // execute a simple Comparison between scalars
        static turbo::Result<type> execute(Datum l, Datum r) {
            if (!l.is_scalar() || !r.is_scalar()) {
                return turbo::invalid_argument_error("Cannot execute Comparison on non-scalars");
            }

            std::vector<Datum> arguments{std::move(l), std::move(r)};

            TURBO_MOVE_OR_RAISE(auto equal, compute::call_function("equal", arguments));

            if (!equal.scalar()->is_valid) return NA;
            if (equal.scalar_as<BooleanScalar>().value) return EQUAL;

            TURBO_MOVE_OR_RAISE(auto less, compute::call_function("less", arguments));

            if (!less.scalar()->is_valid) return NA;
            return less.scalar_as<BooleanScalar>().value ? LESS : GREATER;
        }

        // Given an Expression wrapped in casts which preserve ordering
        // (for example, cast(field_ref("i16"), to_type=int32())), unwrap the inner Expression.
        // This is used to destructure implicitly cast field_refs during Expression
        // simplification.
        static const Expression &StripOrderPreservingCasts(const Expression &expr) {
            auto call = expr.call();
            if (!call) return expr;
            if (call->function_name != "cast") return expr;

            const Expression &from = call->arguments[0];

            auto from_id = from.type()->id();
            auto to_id = expr.type()->id();

            if (is_floating(to_id)) {
                if (is_integer(from_id) || is_floating(from_id)) {
                    return StripOrderPreservingCasts(from);
                }
                return expr;
            }

            if (is_unsigned_integer(to_id)) {
                if (is_unsigned_integer(from_id) && bit_width(to_id) >= bit_width(from_id)) {
                    return StripOrderPreservingCasts(from);
                }
                return expr;
            }

            if (is_signed_integer(to_id)) {
                if (is_integer(from_id) && bit_width(to_id) >= bit_width(from_id)) {
                    return StripOrderPreservingCasts(from);
                }
                return expr;
            }

            return expr;
        }

        static type GetFlipped(type op) {
            switch (op) {
                case NA:
                    return NA;
                case EQUAL:
                    return EQUAL;
                case LESS:
                    return GREATER;
                case GREATER:
                    return LESS;
                case NOT_EQUAL:
                    return NOT_EQUAL;
                case LESS_EQUAL:
                    return GREATER_EQUAL;
                case GREATER_EQUAL:
                    return LESS_EQUAL;
            }
                    DKCHECK(false);
            return NA;
        }

        static std::string GetName(type op) {
            switch (op) {
                case NA:
                    break;
                case EQUAL:
                    return "equal";
                case LESS:
                    return "less";
                case GREATER:
                    return "greater";
                case NOT_EQUAL:
                    return "not_equal";
                case LESS_EQUAL:
                    return "less_equal";
                case GREATER_EQUAL:
                    return "greater_equal";
            }
            return "na";
        }

        static std::string GetOp(type op) {
            switch (op) {
                case NA:
                            DKCHECK(false) << "unreachable";
                    break;
                case EQUAL:
                    return "==";
                case LESS:
                    return "<";
                case GREATER:
                    return ">";
                case NOT_EQUAL:
                    return "!=";
                case LESS_EQUAL:
                    return "<=";
                case GREATER_EQUAL:
                    return ">=";
            }
                    DKCHECK(false);
            return "";
        }
    };

    inline const compute::CastOptions *GetCastOptions(const Expression::Call &call) {
        if (call.function_name != "cast") return nullptr;
        return ::turbo::checked_cast<const compute::CastOptions *>(call.options.get());
    }

    inline bool IsSetLookup(const std::string &function) {
        return function == "is_in" || function == "index_in";
    }

    inline const compute::MakeStructOptions *GetMakeStructOptions(
            const Expression::Call &call) {
        if (call.function_name != "make_struct") return nullptr;
        return ::turbo::checked_cast<const compute::MakeStructOptions *>(
                call.options.get());
    }

    /// A helper for unboxing an Expression composed of associative function calls.
    /// Such expressions can frequently be rearranged to a semantically equivalent
    /// expression for more optimal execution or more straightforward manipulation.
    /// For example, (a + ((b + 3) + 4)) is equivalent to (((4 + 3) + a) + b) and the latter
    /// can be trivially constant-folded to ((7 + a) + b).
    struct FlattenedAssociativeChain {
        /// True if a chain was already a left fold.
        bool was_left_folded = true;

        /// All "branch" expressions in a flattened chain. For example given (a + ((b + 3) + 4))
        /// exprs would be [(a + ((b + 3) + 4)), ((b + 3) + 4), (b + 3)]
        std::vector<Expression> exprs;

        /// All "leaf" expressions in a flattened chain. For example given (a + ((b + 3) + 4))
        /// the fringe would be [a, b, 3, 4]
        std::vector<Expression> fringe;

        explicit FlattenedAssociativeChain(Expression expr) : exprs{std::move(expr)} {
            auto call = CallNotNull(exprs.back());
            fringe = call->arguments;

            auto it = fringe.begin();

            while (it != fringe.end()) {
                auto sub_call = it->call();
                if (!sub_call || sub_call->function_name != call->function_name) {
                    ++it;
                    continue;
                }

                if (it != fringe.begin()) {
                    was_left_folded = false;
                }

                exprs.push_back(std::move(*it));
                it = fringe.erase(it);

                auto index = it - fringe.begin();
                fringe.insert(it, sub_call->arguments.begin(), sub_call->arguments.end());
                it = fringe.begin() + index;
                // NB: no increment so we hit sub_call's first argument next iteration
            }

                    DKCHECK(std::all_of(exprs.begin(), exprs.end(), [](const Expression &expr) {
                return CallNotNull(expr)->options == nullptr;
            }));
        }
    };

    inline turbo::Result<std::shared_ptr<compute::Function>> get_function(
            const Expression::Call &call, compute::ExecContext *exec_context) {
        TURBO_MOVE_OR_RAISE(auto function,
                            exec_context->func_registry()->get_function(call.function_name));
        if (function.get() != exec_context->func_registry()->cast_function()) {
            return function;
        }

        // XXX this special case is strange; why not make "cast" a ScalarFunction?
        const TypeHolder &to_type =
                ::turbo::checked_cast<const compute::CastOptions &>(*call.options).to_type;
        return GetCastFunction(*to_type);
    }

}  // namespace nebula::compute
